package org.usmile.algorithms.learn.tree;

public class RedBlackTree<E extends Comparable<E>> {
    TreeNode<E> root;

    public void add(E e) {
        root = add(root, e);

        root.color = TreeNode.BLACK;
    }

    private TreeNode<E> add(TreeNode<E> node, E e) {
        if (node == null) {
            return new TreeNode<>(e);
        }

        if (e.compareTo(node.e) == 0) {
            return node;
        } else if (e.compareTo(node.e) < 0) {
            node.left = add(node.left, e);
        } else {
            node.right = add(node.right, e);
        }

        return reBalance(node);
    }

    private TreeNode<E> reBalance(TreeNode<E> node) {
        // 维护以 node 为根节点的子树的黑平衡
        if (isRed(node.right) && !isRed(node.left)) {
            node = leftRotate(node);
        }

        if (isRed(node.left) && isRed(node.left.left)) {
            node = rightRotate(node);
        }

        if (isRed(node.left) && isRed(node.right)) {
            flipColors(node);
        }

        return node;
    }

    //    node                    x
    //    /  \       左旋转      /   \
    //   T1   x     ------->  node  T3
    //       / \              /  \
    //      T2 T3            T1  T2
    private TreeNode<E> leftRotate(TreeNode<E> node) {
        TreeNode<E> right = node.right;
        node.right = right.left;
        right.left = node;

        right.color = node.color;
        node.color = TreeNode.RED;

        return right;
    }

    //    node                    x
    //    /  \       右旋转      /   \
    //   x   T2     ------->   y   node
    //  / \                        /  \
    // y  T1                      T1  T2
    private TreeNode<E> rightRotate(TreeNode<E> node) {
        TreeNode<E> left = node.left;
        node.left = left.right;
        left.right = node;

        left.color = node.color;
        node.color = TreeNode.RED;

        return left;
    }

    private void flipColors(TreeNode<E> node) {
        node.color = TreeNode.RED;
        node.left.color = TreeNode.BLACK;
        node.right.color = TreeNode.BLACK;
    }


    private boolean isRed(TreeNode<E> node) {
        if (node == null) {
            return TreeNode.BLACK;
        }

        return node.color;
    }
}
